#!/usr/bin/env python
# -*- coding: utf-8 -*-

'''
Faraday Penetration Test IDE
Copyright (C) 2013  Infobyte LLC (http://www.infobytesec.com/)
See the file 'doc/LICENSE' for the license information

'''
from __future__ import with_statement
import re
import os
import sys
import logging

from faraday.client.plugins import core

logger = logging.getLogger(__name__)

try:
    import xml.etree.cElementTree as ET
    import xml.etree.ElementTree as ET_ORIG
    ETREE_VERSION = ET_ORIG.VERSION
except ImportError:
    import xml.etree.ElementTree as ET
    ETREE_VERSION = ET.VERSION

ETREE_VERSION = [int(i) for i in ETREE_VERSION.split(".")]

current_path = os.path.abspath(os.getcwd())

__author__ = "Francisco Amato"
__copyright__ = "Copyright (c) 2013, Infobyte LLC"
__credits__ = ["Francisco Amato"]
__license__ = ""
__version__ = "1.0.0"
__maintainer__ = "Francisco Amato"
__email__ = "famato@infobytesec.com"
__status__ = "Development"


class ImpactXmlParser(object):
    """
    The objective of this class is to parse an xml file generated by the impact tool.

    TODO: Handle errors.
    TODO: Test impact output version. Handle what happens if the parser doesn't support it.
    TODO: Test cases.

    @param impact_xml_filepath A proper xml generated by impact
    """

    def __init__(self, xml_output):
        tree = self.parse_xml(xml_output)
        if tree:
            self.items = [data for data in self.get_items(tree)]
        else:
            self.items = []

    def parse_xml(self, xml_output):
        """
        Open and parse an xml file.

        TODO: Write custom parser to just read the nodes that we need instead of
        reading the whole file.

        @return xml_tree An xml tree instance. None if error.
        """
        try:
            tree = ET.fromstring(xml_output)
        except SyntaxError as err:
            logger.error("SyntaxError: %s. %s" % (err, xml_output))
            return None

        return tree

    def get_items(self, tree):
        """
        @return items A list of Host instances
        """
        for node in tree.findall("entity/[@class='host']"):
            yield Item(node, tree)


class Item(object):
    """
    An abstract representation of a Item


    @param item_node A item_node taken from an impact xml tree
    """

    def __init__(self, item_node, parent=None):
        self.node = item_node

        self.arch = self.get_text_from_subnode("property/[@key='arch']")

        self.host = self.get_text_from_subnode(
            "property/[@key='display_name']")

        self.ip = self.get_text_from_subnode("property/[@key='ip']")

        self.os = self.get_text_from_subnode(
            "property/[@key='os']/property/[@key='entity name']")

        self.ports = []
        self.services = []
        self.process_ports(item_node)
        self.process_services(item_node)

        self.agent = False

        for node in parent.findall("entity/[@class='agent']"):

            self.node = node
            agentip = node.get('name').split("/")[1]

            if self.ip == agentip:

                self.agentip = agentip

                self.ipfrom = self.get_text_from_subnode(
                    "property/[@key='Connection Properties']/property/[@key='ip']") or agentip

                self.agentype = node.get("type")

                self.agentport = self.get_text_from_subnode(
                    "property/[@key='Connection Properties']//property/[@key='port']") or ""

                self.agentsubtype = self.get_text_from_subnode(
                    "property/[@key='Connection Properties']//property/[@key='subtype']") or ""

                self.agentcon = self.get_text_from_subnode(
                    "property/[@key='Connection Properties']//property/[@key='type']") or ""

                self.agent = True
                break

        self.results = self.getResults(item_node)

    def process_ports(self, item_node):
        for p in item_node.findall("property/[@key='tcp_ports']/property/[@type='port']"):
            self.ports.append({'port': p.get('key'),
                               'protocol': "tcp",
                               'status': "open" if p.text == "listen" else p.text})

        for p in item_node.findall("property/[@key='udp_ports']/property/[@type='port']"):
            self.ports.append({'port': p.get('key'),
                               'protocol': "udp",
                               'status': "open" if p.text == "listen" else p.text})

    def process_services(self, item_node):
        for service in item_node.findall("property/[@key='services']/property"):
            service_name = service.get("key")
            port, protocol = service.findall('property')[0].get('key').split('-')
            self.services.append({
                "name": service_name,
                "protocol": protocol,
                "port": port
            })

    def getResults(self, tree):
        """
        :param tree:
        """
        for self.issues in tree.findall("property/[@key='Vulnerabilities']/property/[@type='container']"):
            yield Results(self.issues)
        # 2017R1 compatibility
        for self.issues in tree.findall("property/[@key='exposures']/property/[@type='container']"):
            yield Results(self.issues)

    def get_text_from_subnode(self, subnode_xpath_expr):
        """
        Finds a subnode in the host node and the retrieves a value from it.

        @return An attribute value
        """
        sub_node = self.node.find(subnode_xpath_expr)
        if sub_node is not None:
            return sub_node.text

        return None


class Results():

    def __init__(self, issue_node):
        self.node = issue_node
        self.ref = [issue_node.get("key")]
        self.severity = ""
        self.port = "Unknown"
        self.service_name = "n/a"
        self.protocol = "tcp?"
        vuln = issue_node.find("property/property")
        if not vuln:
            # 2017R1 compatibility
            self.ref = []
            vuln = issue_node.find("property")
            self.name = self.get_text_from_subnode("property/[@key='title']")
            self.desc = self.get_text_from_subnode("property/[@key='description']")
            self.severity = self.get_text_from_subnode("property/[@key='severity']")
            self.service_name = self.get_text_from_subnode("property/[@key='service']")
        else:
            # 2013R3 xml version
            self.name = vuln.get("key")
            self.node = vuln
            self.desc = self.get_text_from_subnode("property/[@key='description']")
            self.port = self.get_text_from_subnode("property/[@key='port']")

    def get_text_from_subnode(self, subnode_xpath_expr):
        """
        Finds a subnode in the host node and the retrieves a value from it.

        @return An attribute value
        """
        sub_node = self.node.find(subnode_xpath_expr)
        if sub_node is not None:
            return sub_node.text

        return None


class ImpactPlugin(core.PluginBase):
    """
    Example plugin to parse impact output.
    """

    def __init__(self):
        core.PluginBase.__init__(self)
        self.id = "Core Impact"
        self.name = "Core Impact XML Output Plugin"
        self.plugin_version = "0.0.2"
        self.version = "Core Impact 2013R1/2017R2"
        self.framework_version = "1.0.0"
        self.options = None
        self._current_output = None
        self._command_regex = re.compile(r'^(sudo impact|\.\/impact).*?')

        global current_path
        self._output_file_path = os.path.join(self.data_path,
                                              "impact_output-%s.xml" % self._rid)

    def parseOutputString(self, output, debug=False):
        parser = ImpactXmlParser(output)
        mapped_services = {}
        mapped_ports = {}
        for item in parser.items:

            h_id = self.createAndAddHost(
                item.ip,
                item.os + " " + item.arch)

            i_id = self.createAndAddInterface(
                h_id,
                item.ip,
                ipv4_address=item.ip,
                hostname_resolution=[item.host])

            for service in item.services:
                s_id = self.createAndAddServiceToInterface(
                    h_id,
                    i_id,
                    service['name'],
                    service['protocol'],
                    ports=[service['port']],
                    status='open')
                mapped_services[service['name']] = s_id
                mapped_ports[service['port']] = s_id

            if item.agent:
                desc = "Agent Type: " + item.agentype
                desc += "\nConn from:" + item.ipfrom
                desc += "\nPort:" + item.agentport
                desc += "\nProtocol:" + item.agentsubtype
                desc += "\nConn:" + item.agentcon

                self.createAndAddVulnToHost(
                    h_id,
                    "Core Impact Agent",
                    desc=desc,
                    severity="HIGH")

            for v in item.results:
                if v.service_name == "n/a" and v.port == "Unknown":
                    self.createAndAddVulnToHost(
                        h_id,
                        v.name,
                        desc=v.desc,
                        severity=v.severity,
                        ref=v.ref)
                else:
                    s_id = mapped_services.get(v.service_name) or mapped_ports.get(v.port)
                    print(v.service_name)
                    print(s_id)
                    self.createAndAddVulnToService(
                        h_id,
                        s_id,
                        v.name,
                        desc=v.desc,
                        severity=v.severity,
                        ref=v.ref)

            for p in item.ports:
                s_id = self.createAndAddServiceToInterface(
                    h_id,
                    i_id,
                    p['port'],
                    p['protocol'],
                    ports=[p['port']],
                    status=p['status'])
        del parser

    def processCommandString(self, username, current_path, command_string):
        return None

    def setHost(self):
        pass


def createPlugin():
    return ImpactPlugin()


if __name__ == '__main__':
    parser = ImpactXmlParser(sys.argv[1])
    for item in parser.items:
        if item.status == 'up':
            print(item)
